import os
from magic_pdf.data.data_reader_writer import FileBasedDataReader, FileBasedDataWriter
from magic_pdf.data.dataset import PymuDocDataset
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
from magic_pdf.config.enums import SupportedPdfParseMethod

# 1. 准备输入和输出路径
pdf_path      = "input/ITM17-0007.pdf" # PDF 文件
out_image_dir = "output/images"        # 图片输出目录
out_md_dir    = "output/markdown"      # Markdown/JSON 输出目录
os.makedirs(out_image_dir, exist_ok=True)
os.makedirs(out_md_dir,    exist_ok=True)

# 2. 读取 PDF 二进制
reader = FileBasedDataReader("")       # "" 表示本地文件系统
pdf_bytes = reader.read(pdf_path)

# 3. 构建 Dataset 并自动选择 OCR/非 OCR 模式
ds = PymuDocDataset(pdf_bytes)
if ds.classify() == SupportedPdfParseMethod.OCR:
    infer_result = ds.apply(doc_analyze, ocr=True)
    pipe_result  = infer_result.pipe_ocr_mode(FileBasedDataWriter(out_image_dir))
else:
    infer_result = ds.apply(doc_analyze, ocr=False)
    pipe_result  = infer_result.pipe_txt_mode(FileBasedDataWriter(out_image_dir))

# 4. 导出结果
# 4.1 绘制模型预测（可视化校验）
infer_result.draw_model(os.path.join(out_md_dir, "model.pdf"))
# 4.2 绘制布局和文本片段
pipe_result.draw_layout(os.path.join(out_md_dir, "layout.pdf"))
pipe_result.draw_span(os.path.join(out_md_dir,   "spans.pdf"))
# 4.3 获取并保存 Markdown / 中间 JSON
pipe_result.dump_md(FileBasedDataWriter(out_md_dir),
                    "result.md",
                    os.path.basename(out_image_dir))
pipe_result.dump_middle_json(FileBasedDataWriter(out_md_dir),
                             "result_middle.json")